# encoding: utf-8
#!/usr/bin/python3
import os
from torch.utils.data import Dataset,DataLoader
import numpy as np
import pandas as pd
import torch
import json
import cv2
from scipy import ndimage
from data_aug import resize,random_rotate_and_crop,data_aug
from make_heatmap import make_all_heatmap,heatmap_to_anno,make_group_heatmap
import torch.nn.functional as F
import time
from save_heatmap import group_point


def resize_point(image,point,size):
    ori_size = image.shape[:-1]
    image = cv2.resize(image,size[::-1])
    zoom_x = size[0]/ori_size[0]
    zoom_y = size[1]/ori_size[1]
    point = [[p[0]*zoom_x,p[1] * zoom_y] for p in point]
    return image,point

class MyDataset(Dataset):
    def __init__(self,data_path,json_dir,image_size = (768,384),sigma = 15,data_aug = False,Use_Cewei = False,heatmap_path = None):
        ## image_size = {height * width}
        self.data_path = data_path
        self.json_data = json.load(open(json_dir))
        self.data_len = len(self.json_data)
        self.image_size = image_size
        self.singma   = sigma
        self.data_aug = data_aug
        self.Use_Cewei = Use_Cewei
        self.heatmap_path = heatmap_path

    def __getitem__(self, index):
        image = cv2.imread(os.path.join(self.data_path,self.json_data[index]['name'] + '.jpg'))
        
        
        label = self.json_data[index]['point']
        label = [p[0][::-1] for p in label]
        # print('*'*10,len(label))
        if self.image_size is not None:
            image,label  = resize_point(image,label,self.image_size)
        start_time = time.time()
        if self.data_aug:
            try:
                image,label = data_aug(image,label)
            except Exception:
                pass
        # print('-'*10,len(label))
        # print('use time in aug is --',time.time() - start_time)
        image_shape = image.shape
        ## 每个坐标值对应一层heatmap
        start_time = time.time()
        label = 255*make_all_heatmap(image_shape[1],image_shape[0],anno=label,sigma=self.singma)
        
        
        # ## 分组生成高斯热图
        # label = group_point(label)
        # label = 255*make_group_heatmap(image_shape[1],image_shape[0],anno=label,sigma=self.singma)
        # start_time = time.time()
        
        # label = np.load(os.path.join(self.heatmap_path,self.json_data[index]['name'] + '.npy'))
        # print('the time in load heatmap is --',time.time() - start_time)
        #
        # # 将heatmap进行可视化操作，用来检查
        # show_label = (np.sum(label,axis=0)).astype(np.uint8)
        # cv2.imshow('heatmap',show_label)
        # cv2.waitKey()


        # label = make_heatmap_z_c(image_shape[1], image_shape[0], anno=label, sigma=self.singma,zheng=False)
        
        # io.imshow(label[1]+label[0])
        # plt.show()
        # 生成heatmap之后。在进行随机的旋转和裁剪操作。
        # print('before rotate and crop', image.shape, label.shape)
        
        # start_time = time.time()
        image,label = random_rotate_and_crop(image,label)
        # print('the time in crop and rotate data is --',time.time() - start_time)
        
        
        
        # print('after rotate and crop',image.shape,label.shape)
        ## resize 到统一的尺寸用于训练。
        # start_time = time.time()
        size = image.shape
        image = ndimage.zoom(image,(self.image_size[0]/size[0],self.image_size[1]/size[1],1),order=2)
        label = ndimage.zoom(label,(1,self.image_size[0]/size[0],self.image_size[1]/size[1]),order=2)
        # print('the time in zoom data is --',time.time() - start_time)
        
        
        
        # cv2.imshow('heatmap',label[0].astype(np.uint8))
        # cv2.waitKey(0)
        
        image = torch.from_numpy(np.transpose(image,axes=(2,0,1))).float()
        label = torch.from_numpy(label)
        # print(label)
        # print('image shape', image.shape, 'label shape', label.shape)
        return image,label

    def __len__(self):
        return self.data_len

if __name__ == '__main__':
    Data_Path = './data/image'
    Json_Dir = 'label_file/new_point_train.json'
    Image_Size = (640, 512)
    Batch_Size = 8
    Sigma = [9,9,17,25]

    My_dataSet = MyDataset(data_path=Data_Path, json_dir=Json_Dir, image_size=Image_Size,
                           data_aug=True, sigma=Sigma,heatmap_path='./data/heatmap_640_512')
    train_data_loader = DataLoader(dataset=My_dataSet, batch_size=Batch_Size, shuffle=True,
                                   num_workers=0, pin_memory=True)
    for batch_image,batch_label in train_data_loader:
        pass
        print(batch_image.shape,batch_label.shape)